# -*- coding: utf-8 -*-
# @Function: LLM查询测试器
# @Description: 使用大模型回答测试问题，结合数据库查询和LLM生成能力，提供更智能的问答功能
# @Usage: 可独立运行或被其他模块导入，提供LLM增强的问答功能
# @Dependencies: db_manager.py, tokenizer.py, 外部LLM服务

import sys
import os

# 将项目根目录添加到 Python 路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, project_root)

import os
import json
import re
from typing import List, Dict
from dotenv import load_dotenv
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from D_db.core.db_manager import DBManager
from D_db.core.tokenizer import Tokenizer
from D_db.tools.prompt_utils import build_llm_context_from_qa

# 加载环境变量
load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

class LLMQueryTester:
    def __init__(self):
        """初始化LLM查询测试器"""
        self.db_manager = DBManager()
        self.tokenizer = Tokenizer()
        
        # 初始化阿里云百炼模型
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7
        )
        
        # 创建回答生成的提示模板
        self.answer_template = PromptTemplate(
            input_variables=["question", "context"],
            template="""你是一个专业的水库调度规程问答助手。请根据以下参考信息，回答用户的问题。

用户问题：
{question}

参考信息：
{context}

要求：
1. 基于参考信息提供准确、完整的回答
2. 如果参考信息不足以回答问题，请明确指出
3. 回答要简洁明了，直接针对问题给出答案
4. 不要在回答中提及"根据参考信息"等字样
5. 使用专业、客观的语气
6. 请在答案末尾标注信息来源：
   - 如果使用了表格内容，格式为：[table_name: 实际表格名称]
     例如：[table_name: 防洪限制水位]
   - 如果使用了文本内容，格式为：[source: 章节路径]
     例如：[source: 防洪调度/调度原则]

请提供你的回答：
"""
        )
        
        # 创建回答生成链 (使用新的 RunnableSequence 语法)
        self.answer_chain = self.answer_template | self.llm

    def test_query(self, test_question: str, limit: int = 20):
        """测试查询功能并使用大模型生成回答
    
        Args:
            test_question (str): 测试问题
            limit (int, optional): 返回结果数量限制. Defaults to 20.
    
        Returns:
            dict: 查询结果和大模型回答
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 对测试问题进行分词
            token_data = self.tokenizer.tokenize_question(test_question)
            tokens = token_data.get('tokens', [])
            category_id = token_data.get('category_id')  # 获取提取的category_id
            
            if not tokens:
                print("测试问题分词结果为空")
                return {"answer": "无法处理您的问题，请尝试重新表述。"}
                
            # 查询相关问题，传入category_id参数
            results = self.db_manager.query_qa(tokens, limit, category_id)
            
            # 输出查询结果
            print(f"\n测试问题: {test_question}")
            print(f"分词结果: {tokens}")
            if category_id:
                print(f"识别的水库ID: {category_id}")
            print(f"查询到 {len(results)} 个相关问题")
            
            # 复用共享工具模块构造 LLM 上下文和结构化 QA 结果
            context_for_llm, qa_results = build_llm_context_from_qa(results)

            # 打印每条结果的摘要（与之前类似）
            for i, item in enumerate(qa_results, 1):
                print(f"结果 {i}: {item['question']} (类别: {item.get('category_id')}, 匹配度: {item.get('match_count')})")
                print(f"答案 {i}: {item['answer']}")
                if item.get('context'):
                    print(f"上下文: {item['context']}")
                print()

            # 将最终提示写入日志，便于检查
            try:
                log_path = os.path.join(project_root, "context_for_llm.log")
                with open(log_path, "w", encoding="utf-8") as f:
                    f.write(context_for_llm)
            except Exception:
                pass

            # 如果没有查询结果，返回提示信息
            if not results:
                return {"answer": "抱歉，没有找到与您问题相关的信息。"}

            # 使用大模型生成回答
            print("\n使用大模型生成回答...")
            # llm_answer = self.answer_chain.run(question=test_question, context=context_for_llm)
            llm_answer = self.answer_chain.invoke({"question": test_question, "context": context_for_llm})
            print(f"大模型回答: {llm_answer.content}")
            final_answer = llm_answer.content

            # 检测回答中是否包含表格标记和文本来源标记
            try:
                # 1. 处理表格标记
                table_patterns = [
                    re.compile(r"\[table_name:\s*([^\]]+)\]"),  # [table_name: xxx]
                    re.compile(r"\[表[0-9\-]*([^\]]+)\]"),       # [表3-4-1防洪限制水位]
                ]
                
                table_images_map = {}
                found_tables = []
                all_markers = []
                
                for pattern in table_patterns:
                    matches = pattern.findall(final_answer)
                    if matches:
                        found_tables.extend(matches)
                        all_markers.extend(pattern.finditer(final_answer))
                
                # 从答案中移除表格标记
                if all_markers:
                    all_markers = sorted(set(m.group(0) for m in all_markers), 
                                        key=lambda x: final_answer.rfind(x), reverse=True)
                    for marker in all_markers:
                        final_answer = final_answer.replace(marker, '', 1)
                    final_answer = final_answer.strip()
                
                if found_tables:
                    unique_tables = []
                    for t in found_tables:
                        t = t.strip()
                        if t and t not in unique_tables:
                            unique_tables.append(t)
                    
                    print(f"🔍 检测到表格标记: {unique_tables}")
                    
                    for table_name in unique_tables:
                        try:
                            images = self.db_manager.get_table_images(table_name)
                            urls = []
                            for img in images:
                                if isinstance(img, dict):
                                    url = img.get('url') or img.get('filename')
                                    if url:
                                        urls.append(url)
                                else:
                                    urls.append(str(img))

                            if urls:
                                table_images_map[table_name] = urls
                                print(f"✅ 找到表格图片: {table_name} ({len(urls)} 张)")
                            else:
                                print(f"⚠️ 表格 '{table_name}' 在数据库中没有图片")
                        except Exception as ex_table:
                            print(f"❌ 查询表格图片失败 ({table_name}): {ex_table}")
                
                # 2. 处理文本来源标记 [source: xxx]
                source_pattern = re.compile(r"\[source:\s*([^\]]+)\]")
                text_sources = source_pattern.findall(final_answer)
                
                # 从答案中移除文本来源标记
                if text_sources:
                    final_answer = source_pattern.sub('', final_answer).strip()
                    print(f"📖 检测到文本来源: {text_sources}")
                
                # 如果大模型没有返回来源标记，则从 qa_results 中提取 context
                if not text_sources and not found_tables and qa_results:
                    # 提取前3个QA结果的context作为文本来源
                    for qa in qa_results[:3]:
                        ctx = qa.get('context', '')
                        if ctx and not qa.get('is_table'):
                            # 尝试解析context中的章节信息
                            try:
                                ctx_json = json.loads(ctx) if isinstance(ctx, str) else ctx
                                if isinstance(ctx_json, dict):
                                    h1 = ctx_json.get('h1', '')
                                    h2 = ctx_json.get('h2', '')
                                    if h1 or h2:
                                        source = f"{h1}/{h2}" if h1 and h2 else (h1 or h2)
                                        if source and source not in text_sources:
                                            text_sources.append(source)
                            except:
                                pass
                    
                    if text_sources:
                        print(f"📖 从context提取文本来源: {text_sources}")
                        
            except Exception as ex:
                print(f"处理来源标记时出错: {ex}")

            return {
                "question": test_question,
                "answer": final_answer,
                "qa_results": qa_results,
                "table_images": table_images_map,
                "text_sources": text_sources if 'text_sources' in locals() else []
            }
        except Exception as e:
            print(f"查询失败: {e}")
            return {"answer": f"处理问题时出错: {str(e)}"}
        finally:
            # 关闭数据库连接
            self.db_manager.close()

def main():
    """主函数"""
    print("=== LLM问答测试工具 ===")
    
    # 验证API密钥
    if not DASHSCOPE_CONFIG["api_key"]:
        print("错误: 请设置DASHSCOPE_API_KEY环境变量")
        return
    
    print(f"连接到阿里云百炼服务: {DASHSCOPE_CONFIG['base_url']}")
    print(f"使用模型: {DASHSCOPE_CONFIG['model']}")
    
    tester = LLMQueryTester()
    
    # 测试问题列表
    test_questions = [
        "大伙房水库的面积是多少？"
    ]
    
    # 交互式问答
    while True:
        print("\n" + "=" * 50)
        
        choice = "1"
        
        if choice == "0":
            break
        elif choice == "1":
            for i, question in enumerate(test_questions, 1):
                print(f"\n测试问题 {i}: {question}")
                print("-" * 50)
                result = tester.test_query(question)
                # 不再重复输出大模型回答
                print("=" * 50)
            break
        elif choice == "2":
            question = input("\n请输入您的问题: ")
            if question.strip():
                print("-" * 50)
                result = tester.test_query(question)
                # 不再重复输出大模型回答
                print("=" * 50)
        else:
            print("无效的选择，请重新输入")

if __name__ == "__main__":
    main()